'''
@Author: your name
@Date: 2020-06-02 15:33:30
LastEditTime: 2020-09-18 10:09:07
LastEditors: Please set LastEditors
@Description: In User Settings Edit
@FilePath: /mtl-text-recognition/config.py
'''
# coding=utf-8

import os
import torch


class ConfigOpt:
    def __init__(self):
        self.cur_path = os.path.abspath(os.path.dirname(__file__))
        self.workers = 2
        self.batch_size = 128
        self.saved_model = os.path.join(self.cur_path, "saved_models/8-24_NoJYFW_YYZZ_train-AT-TMixCrop/mtl_best_accuracy.pth")
        self.batch_max_length = 50
        self.imgH = 32
        self.imgW = 512
        self.rgb = None
        self.sensitive = True
        self.PAD = True
        self.Transformation = 'None'  # None|TPS
        self.FeatureExtraction = 'VGG'  # VGG|RCNN|ResNet|MobileNetV3|DenseNet|CNN_Lite
        self.SequenceModeling = 'BiLSTM'  # None|BiLSTM
        self.Prediction = 'CTC'  # CTC|Attn
        self.num_fiducial = 20
        self.input_channel = 1

        self.output_channel = 512
        # self.output_channel = 768
        # self.output_channel = 480
        self.hidden_size = 256
        
        # self.hidden_size = 384
        self.num_gpu = torch.cuda.device_count()
        # self.char_dict = "config/dict/cleaned_all_dict.txt"
        # self.char_dict = 'config/dict/dcp_dict.txt'
        # self.char_dict = 'config/dict/dict.txt'
        self.char_dict = "config/dict/DCP_all_dict.txt"
        self.character = self.get_character()
        self.mtl = True
        self.ctc_num_class = 0
        self.num_class = 0


        '''SRN'''
        self.position_dim = 26
        self.batch_max_character = 25

    def get_character(self):
        ch_chars = ""
        ch_path = os.path.join(self.cur_path, self.char_dict)
        with open(ch_path) as charf:
            for line in charf:
                line = line.strip()
                ch_chars += line.encode("utf-8", 'strict').decode("utf-8", 'strict')
        return ch_chars
